import numpy as np
import pandas as pd
import time

# 随机种子
np.random.seed(2)

# 状态个数
N_STATES = 6
# 动作
ACTIONS = ['left', 'right']
# 策略概率
EPSILON = 1
# 学习比率
ALPHA = 0.1
# 折损因子
LAMBDA = 0.9
# 最大回合（运行次数）
MAX_EPISODES = 30
# 刷新世纪N
REFRESH_TIME = 0.1


def build_q_table(n_states, actions):
    table = pd.DataFrame(np.zeros((n_states, len(actions))), columns=actions)
    return table


def choose_action(state, q_table):
    # 取该状态的所有行动
    state_actions = q_table.iloc[state, :]
    # 随机
    if (np.random.uniform() > EPSILON) or (state_actions.all() == 0):
        action_name = np.random.choice(ACTIONS)
    # 策略
    else:
        # 取最大值
        action_name = state_actions.idxmax()
    return action_name


def get_env_feedback(S, A):
    # 右移
    if A == 'right':
        # 已达最大状态
        if S == N_STATES - 2:
            S = 'terminal'
            R = 1
        else:
            S = S + 1
            R = 0
    # 左移
    else:
        R = 0
        if S > 0:
            S = S - 1
    return S, R


def update_env(S, episode, step_counter):
    env_list = ['-'] * (N_STATES - 1) + ['T']
    if S == 'terminal':
        interaction = '\nEpisode %s: total_steps = %s' % (episode + 1, step_counter)
        print('{}'.format(interaction))
        time.sleep(1)
        print('\r                               ', end='')
    else:
        env_list[S] = 'o'
        interaction = ''.join(env_list)
        print('\r{}'.format(interaction), end='')
        time.sleep(REFRESH_TIME)


def rl():
    q_table = build_q_table(N_STATES, ACTIONS)
    for episode in range(MAX_EPISODES):
        # 初始化
        step_counter = 0
        S = 0
        is_terminated = False
        update_env(S, episode, step_counter)
        # 没有结束就一直循环
        while not is_terminated:
            A = choose_action(S, q_table)
            S_new, R = get_env_feedback(S, A)
            # 取该状态下的该行动的q值
            q_predict = q_table.loc[S, A]
            # 没有胜利
            if S_new != 'terminal':
                q_target = R + LAMBDA * q_table.iloc[S_new, :].max()
            # 胜利
            else:
                q_target = R
                is_terminated = True
            # 更新Q表
            q_table.loc[S, A] += ALPHA * (q_target - q_predict)
            # 更新状态
            S = S_new
            # 更新步数
            step_counter += 1
            # 更新环境
            update_env(S, episode, step_counter)
        # 一回合结束
        # print(q_table)
    return q_table


if __name__ == "__main__":
    q_table = rl()
